import asyncio

from starlette.responses import PlainTextResponse

from starlette.status import HTTP_408_REQUEST_TIMEOUT

from starlette.middleware.base import RequestResponseEndpoint, BaseHTTPMiddleware, DispatchFunction

from starlette.types import ASGIApp


class TimeoutMiddleware(BaseHTTPMiddleware):
    '''
    需注意，一般基于BaseHTTPMiddleware胡中间件建议都放置在最前面进行处理
    '''
    def __init__(self, app: ASGIApp, timeout: int = 15):
        super().__init__(app)
        self.timeout = timeout

    async def dispatch(self, request, call_next):
        try:
            return await asyncio.wait_for(call_next(request), timeout=self.timeout)
        except asyncio.TimeoutError:
            return PlainTextResponse(status_code=HTTP_408_REQUEST_TIMEOUT, content='timeout')
